'''
@author: morales
Updated Jan 12, 2010
'''

from easygui import *
from Parameters import *
import math

class Dialog:
	
	fieldValues = None
	
	def getResponse(self):
		return None
	
	def getMEBResponse(self, msg, title, fieldNames, defaultValues ):
		''' Used to get an array of responses from a ME Box '''
		
		fieldValues = multenterbox( msg,title, fieldNames, defaultValues )
		
		# make sure that none of the fields was left blank
		while 1:
			if fieldValues == None: break
			errmsg = ""
			for i in range(len(fieldNames)):
				if fieldValues[i].strip() == "":
					errmsg = errmsg + ('"%s" is a required field.\n\n' % fieldNames[i])
			if errmsg == "": break # no problems found
			fieldValues = multenterbox(errmsg, title, fieldNames, fieldValues)
		
		self.fieldValues = fieldValues
		return self.fieldValues
	
	def __str__(self):
		print self.fieldValues

'''
Used to choose which Demography to explore
Updated Jan 12, 2010
'''
class Dialog_MEMES(Dialog):

	def __init__(self):
		self.memes = ["C: Wright-Fisher", "G: Growth", \
					"CC: Constant-Constant", "CG: Constant to Exp growth"]
			
	def getResponse(self):	
		self.response = buttonbox("Choose the MEMES of interest:", "MEMES Type", \
		    [ x for x in self.memes ] )
		return self.response
	
	def __str__(self):
		return self.response

'''
Input population parameters
Updated Jan 12, 2010
'''
class Dialog_Population(Dialog):

	def __init__(self):
		
		self.popParameters = Population()
		self.setDefaultPopulationParameters()
		
	def setDefaultPopulationParameters(self):
		
		self.popParameters.numChromosomes = 30
		self.popParameters.numSegregatingSites = 50

		self.popParameters.rateMutation = 1E-8
		self.popParameters.rateRecombination = 0.0
		self.popParameters.sequenceLength = 25000

		self.popParameters.N = 10000
		self.popParameters.numSimulations = 50000
		self.popParameters.numPowerSimulations = 5000
		
	def getResponse(self):
		
		msg = "Enter parameters below"
		title = "Population parameters"
		
		fieldNames = [ \
			"n - Number of Chromosomes", \
			"S - Segregating sites", \
			"N(0) - Present effective size", \
			"Number of simulations", \
			"Power simulations amount"]
		
		defaultValues = [ \
			str(self.popParameters.numChromosomes), \
			str(self.popParameters.numSegregatingSites), \
			str(self.popParameters.N), \
			str(self.popParameters.numSimulations), \
			str(self.popParameters.numPowerSimulations)]

		fieldValues = self.getMEBResponse( msg, title, fieldNames, defaultValues )

		self.popParameters.numChromosomes = int( fieldValues[0].strip() )
		self.popParameters.numSegregatingSites = int( fieldValues[1].strip() )
		self.popParameters.N = int( fieldValues[2].strip() )
		self.popParameters.numSimulations = int( fieldValues[3].strip() )
		self.popParameters.numPowerSimulations = int( fieldValues[4].strip() )
		
		return self.popParameters

	def __str__(self):
		return str(self.popParameters)
'''
Get the demographic parameters
'''
class Dialog_Demography(Dialog):
	
	def __init__(self, N, numEpochs):
		self.N = N
		self.numEpochs = numEpochs
		
		self.alpha = []
		self.tau = []
		self.isConstant = []
		
		self.demography = Demography(self.N)

	def updateDemography(self):

		self.demography = Demography(self.N)
		
		for i in range( 0, self.numEpochs ):
			self.demography.addEpoch( self.alpha[i], self.tau[i], self.isConstant[i] )
			
	def __str__(self):
		return str(self.demography)


class Dialog_Demography_C( Dialog_Demography ):
	
	def __init__(self, N):
		Dialog_Demography.__init__(self, N, 1)
		
		self.alpha.append(N)
		self.tau.append(0.0)
		self.isConstant.append(True)
		
	def getResponse(self):
		self.updateDemography()
		return self.demography
			

class Dialog_Demography_G( Dialog_Demography ):
	
	def __init__(self, N):
		Dialog_Demography.__init__(self, N, 1)
		
		self.alpha.append(20.0)
		self.tau.append(0.0)
		self.isConstant.append(False)

	def getResponse(self):
		
		msg = "Enter G parameters (time in generations) where:\n"
		msg += "  N = %i at\n  time = 0 gen" % (self.N)
		
		title = "G-model parameters"
		
		fieldNames = ["alpha - Growth/decay rate"]
		defaultValues = [ str(self.alpha[0]) ]
						
		fieldValues = self.getMEBResponse( msg,title, fieldNames, defaultValues )

		self.alpha[0] = float( fieldValues[0].strip() )
		
		self.updateDemography()
		return self.demography
				

class Dialog_Demography_CC( Dialog_Demography ):

	def __init__(self, N):
		Dialog_Demography.__init__(self, N, 2)
		
		self.alpha.append(N)
		self.tau.append(0)
		self.isConstant.append(True)
		
		self.alpha.append(N)
		self.tau.append(0)
		self.isConstant.append(True)
		
	def getResponse(self):
		
		msg = "Enter CC parameters (time in generations) where:\n"
		msg += "  N = %i at\n  time = 0 gen" % (self.N)
		
		title = "CC-model parameters"
		
		fieldNames = ["N1 - New population size", "t1 - Time of population change"]
		defaultValues = [ str(self.alpha[1]), str(self.tau[1]) ]
						
		fieldValues = self.getMEBResponse( msg,title, fieldNames, defaultValues )

		self.alpha[1] = int( fieldValues[0].strip() )
		self.tau[1] = int( fieldValues[1].strip() )
		
		self.updateDemography()
		return self.demography


class Dialog_Demography_CG( Dialog_Demography ):

	def __init__(self, N):
		Dialog_Demography.__init__( self, N, 2 )
		
		self.alpha.append(10.0)
		self.tau.append(0)
		self.isConstant.append(False)
		
		self.alpha.append(0.0)
		self.tau.append(1000)
		self.isConstant.append(True)
		
	def getResponse(self):
		
		msg = "Enter CG parameters (time in generations) where:\n"
		msg += "  N = %i at\n  time = 0 gen" % (self.N)
		
		title = "CG-model parameters"
		
		fieldNames = ["alpha0 - Growth rate", "t1 - Time of population change"]
		defaultValues = [ str(self.alpha[0]), str(self.tau[1]) ]
						
		fieldValues = self.getMEBResponse( msg,title, fieldNames, defaultValues )

		a = self.alpha[0] = float( fieldValues[0].strip() )
		t = self.tau[1] = 1.0 * int( fieldValues[1].strip() )
		
		
		self.alpha[1] = 1.0 * self.N * math.exp(-a * t / (4.0 * self.N))
		
		self.updateDemography()
		return self.demography
				